// 针对网络模块的压力测试
#include <iostream>
#include <string>
#include <cstring>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <unistd.h>
#include <thread>
#include <chrono>
#include <cstdlib>
#include <atomic>
#include <ctime>
#include <vector>
#include <mutex>
#include <iomanip>
#include "include/net/crc32.h"
#include <algorithm>
using namespace std;

// 配置参数
int THREAD_COUNT = 200;                 // 并发线程数
int PACKETS_PER_CONNECTION = 1000;      // 每个连接发送的包数
int packSize = 0;                       // 包大小1024倍数
const int SERVER_PORT = 11111;          // 服务器端口
const char *SERVER_IP = "127.0.0.1";    // 服务器IP
const int CONNECTION_TIMEOUT_MS = 2000; // 连接超时(毫秒)
const int RECV_TIMEOUT_MS = 100;        // 接收超时(毫秒)

// 性能统计
atomic<int> g_connSuccess(0);                          // 成功连接数
atomic<int> g_connFailed(0);                           // 失败连接数
atomic<long> g_totalPacketsSent(0);                    // 发送的总包数
atomic<long> g_totalBytesSent(0);                      // 发送的总字节数
atomic<long> g_totalPacketsReceived(0);                // 接收的总包数
atomic<long> g_totalBytesReceived(0);                  // 接收的总字节数
atomic<long> g_totalProcessingTime(0);                 // 总处理时间(微秒)
chrono::high_resolution_clock::time_point g_startTime; // 测试开始时间

// 延迟统计
vector<long> g_latencyRecords; // 存储每个请求的延迟
mutex g_latencyMutex;          // 保护延迟记录

// 连接错误信息统计
mutex g_errorMutex;
vector<string> g_errorMessages;

// 包头结构体
#pragma pack(1) // 设置1字节对齐
typedef struct _COMM_PKG_HEADER
{
    unsigned short pkgLen;  // 报文总长度（2字节）包头+包体
    unsigned short msgCode; // 消息类型（2字节）
    int crc32;              // CRC32校验码（4字节）
} COMM_PKG_HEADER, *LPCOMM_PKG_HEADER;
#pragma pack() // 恢复默认对齐

// 设置socket超时
bool SetSocketTimeout(int sockfd, int type, int timeout_ms)
{
    struct timeval tv;
    tv.tv_sec = timeout_ms / 1000;
    tv.tv_usec = (timeout_ms % 1000) * 1000;
    return setsockopt(sockfd, SOL_SOCKET, type, &tv, sizeof(tv)) == 0;
}

// 单线程处理函数
void ThreadFunc()
{
    // 创建socket
    int sockfd = socket(AF_INET, SOCK_STREAM, 0);
    if (sockfd < 0)
    {
        lock_guard<mutex> lock(g_errorMutex);
        g_errorMessages.push_back("压力测试套接字生成失败: " + string(strerror(errno)));
        g_connFailed++;
        return;
    }

    // 设置连接超时
    if (!SetSocketTimeout(sockfd, SO_SNDTIMEO, CONNECTION_TIMEOUT_MS))
    {
        lock_guard<mutex> lock(g_errorMutex);
        g_errorMessages.push_back("压力测试设置连接超时失败: " + string(strerror(errno)));
        g_connFailed++;
        close(sockfd);
        return;
    }

    // 配置服务器地址
    struct sockaddr_in serverAddr;
    memset(&serverAddr, 0, sizeof(serverAddr));
    serverAddr.sin_family = AF_INET;
    serverAddr.sin_port = htons(SERVER_PORT);

    if (inet_pton(AF_INET, SERVER_IP, &serverAddr.sin_addr) <= 0)
    {
        lock_guard<mutex> lock(g_errorMutex);
        g_errorMessages.push_back("压力测试配置服务器地址失败: " + string(SERVER_IP));
        g_connFailed++;
        close(sockfd);
        return;
    }

    // 连接服务器
    if (connect(sockfd, (struct sockaddr *)&serverAddr, sizeof(serverAddr)) < 0)
    {
        lock_guard<mutex> lock(g_errorMutex);
        g_errorMessages.push_back("压力测试连接失败: " + string(strerror(errno)));
        g_connFailed++;
        close(sockfd);
        return;
    }

    // 连接成功后设置接收超时
    if (!SetSocketTimeout(sockfd, SO_RCVTIMEO, RECV_TIMEOUT_MS))
    {
        lock_guard<mutex> lock(g_errorMutex);
        g_errorMessages.push_back("设置接收超时失败: " + string(strerror(errno)));
    }

    // 连接成功
    g_connSuccess++;

    // 准备发送数据包
    int len = packSize * 1024;
    size_t headLen = sizeof(COMM_PKG_HEADER);
    const size_t total_len = headLen + len;
    char buff[total_len];
    LPCOMM_PKG_HEADER header = reinterpret_cast<LPCOMM_PKG_HEADER>(buff);
    if (packSize != 0)
    {
        string body(packSize * 1024, 'A');
        CCRC32 *m_crc = CCRC32::GetInstance();
        std::vector<unsigned char> tmp_buf(body.begin(), body.end());
        int calccrc = m_crc->Get_CRC(tmp_buf.data(), tmp_buf.size());
        header->crc32 = htonl(calccrc);
    }
    else
    {
        header->crc32 = htonl(0); // 空校验码
    }
    header->pkgLen = htons(static_cast<unsigned short>(total_len));
    header->msgCode = htons(6); // 示例消息类型

    // 接收缓冲区
    char recvBuff[sizeof(COMM_PKG_HEADER)];

    for (int i = 0; i < PACKETS_PER_CONNECTION; i++)
    {
        auto sendStart = chrono::high_resolution_clock::now();

        // 发送数据包
        ssize_t bytesSent = send(sockfd, buff, total_len, MSG_NOSIGNAL);
        if (bytesSent < 0)
        {
            lock_guard<mutex> lock(g_errorMutex);
            g_errorMessages.push_back("压力测试发送失败: " + string(strerror(errno)));
            break;
        }
        g_totalPacketsSent++;
        g_totalBytesSent += bytesSent;

        // 接收服务器响应
        ssize_t bytesReceived = recv(sockfd, recvBuff, sizeof(recvBuff), 0);
        auto recvEnd = chrono::high_resolution_clock::now();

        // 计算处理延迟(微秒)
        auto latency = chrono::duration_cast<chrono::microseconds>(recvEnd - sendStart).count();
        g_totalProcessingTime += latency;

        {
            lock_guard<mutex> lock(g_latencyMutex);
            g_latencyRecords.push_back(latency);
        }

        if (bytesReceived > 0)
        {
            g_totalPacketsReceived++;
            g_totalBytesReceived += bytesReceived;
        }
        else if (bytesReceived == 0)
        {
            // 连接关闭
            lock_guard<mutex> lock(g_errorMutex);
            g_errorMessages.push_back("服务器关闭连接");
            break;
        }
        else if (errno != EAGAIN && errno != EWOULDBLOCK)
        {
            // 真正的接收错误
            lock_guard<mutex> lock(g_errorMutex);
            g_errorMessages.push_back("接收错误: " + string(strerror(errno)));
        }
    }

    close(sockfd);
}

// 计算百分位延迟
long CalculatePercentile(int percentile)
{
    if (g_latencyRecords.empty())
        return 0;

    lock_guard<mutex> lock(g_latencyMutex);
    vector<long> sorted = g_latencyRecords;
    sort(sorted.begin(), sorted.end());

    int index = (percentile * sorted.size()) / 100;
    if (index >= sorted.size())
        index = sorted.size() - 1;

    return sorted[index];
}

// 打印统计信息
void PrintStatistics()
{
    auto endTime = chrono::high_resolution_clock::now();
    auto duration = chrono::duration_cast<chrono::milliseconds>(endTime - g_startTime).count();
    double durationSec = duration / 1000.0;

    // 计算延迟统计
    long minLatency = 0, maxLatency = 0, avgLatency = 0;
    long p50 = 0, p90 = 0, p99 = 0;

    if (!g_latencyRecords.empty())
    {
        minLatency = *min_element(g_latencyRecords.begin(), g_latencyRecords.end());
        maxLatency = *max_element(g_latencyRecords.begin(), g_latencyRecords.end());
        avgLatency = g_totalProcessingTime / g_latencyRecords.size();
        p50 = CalculatePercentile(50);
        p90 = CalculatePercentile(90);
        p99 = CalculatePercentile(99);
    }

    cout << "\n======= 压力测试结果 =======" << endl;
    cout << "服务器地址:        " << SERVER_IP << ":" << SERVER_PORT << endl;
    cout << "并发线程数:        " << THREAD_COUNT << endl;
    cout << "总连接尝试:        " << (g_connSuccess + g_connFailed) << endl;
    cout << "成功连接数:        " << g_connSuccess << endl;
    cout << "失败连接数:        " << g_connFailed << endl;
    cout << "总发送包数:        " << g_totalPacketsSent << endl;
    cout << "总接收包数:        " << g_totalPacketsReceived << endl;
    cout << "发送字节总数:      " << g_totalBytesSent << " bytes" << endl;
    cout << "接收字节总数:      " << g_totalBytesReceived << " bytes" << endl;
    cout << "测试持续时间:      " << duration << " ms (" << fixed << setprecision(2) << durationSec << " sec)" << endl;
    cout << "连接成功率:        " << fixed << setprecision(2)
         << (g_connSuccess * 100.0 / (g_connSuccess + g_connFailed)) << "%" << endl;
    cout << "平均包发送速率:    " << fixed << setprecision(2)
         << (g_totalPacketsSent / durationSec) << " packets/sec" << endl;
    cout << "平均包接收速率:    " << fixed << setprecision(2)
         << (g_totalPacketsReceived / durationSec) << " packets/sec" << endl;
    cout << "平均数据吞吐量:    " << fixed << setprecision(2)
         << ((g_totalBytesSent + g_totalBytesReceived) / (1024.0 * 1024.0) / durationSec) << " MB/sec" << endl;
    cout << "------ 服务器性能指标 ------" << endl;
    cout << "平均延迟:          " << avgLatency << " μs" << endl;
    cout << "最小延迟:          " << minLatency << " μs" << endl;
    cout << "最大延迟:          " << maxLatency << " μs" << endl;
    cout << "P50延迟:           " << p50 << " μs" << endl;
    cout << "P90延迟:           " << p90 << " μs" << endl;
    cout << "P99延迟:           " << p99 << " μs" << endl;
    cout << "请求成功率:        " << fixed << setprecision(2)
         << (g_totalPacketsReceived * 100.0 / g_totalPacketsSent) << "%" << endl;
    cout << "===========================" << endl;

    if (!g_errorMessages.empty())
    {
        cout << "\n====== 错误统计 (" << g_errorMessages.size() << ") ======" << endl;
        for (const auto &msg : g_errorMessages)
        {
            cout << " • " << msg << endl;
        }
    }
}

int main()
{
    cout << "开始压力测试..." << endl;
    cout << "请输入线程个数:";
    cin >> THREAD_COUNT;
    cout << "请输入单线程发送的包数:";
    cin >> PACKETS_PER_CONNECTION;
    cout << "请输入包大小(1024倍数):";
    cin >> packSize;

    g_startTime = chrono::high_resolution_clock::now();
    vector<thread> threads;
    threads.reserve(THREAD_COUNT);

    // 创建线程
    for (int i = 0; i < THREAD_COUNT; i++)
    {
        threads.emplace_back(ThreadFunc);
    }

    // 等待所有线程完成
    for (auto &t : threads)
    {
        t.join();
    }

    // 打印统计结果
    PrintStatistics();

    return 0;
}